import sys, os
import cv2
import argparse
import numpy as np
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PIL import Image, ImageEnhance
from shape import Shape

CURSOR_DEFAULT = Qt.ArrowCursor
CURSOR_POINT = Qt.PointingHandCursor
CURSOR_DRAW = Qt.CrossCursor
CURSOR_MOVE = Qt.ClosedHandCursor
CURSOR_GRAB = Qt.OpenHandCursor

class ImageLabel(QLabel):
    CREATE, EDIT = 0, 1

    croped = pyqtSignal()
    resized = pyqtSignal()
    def __init__(self, parent=None, createMode='rectangle'):
        super(ImageLabel, self).__init__()
        self.parent = parent
        self.image = None
        self.current = None
        self.imgW, self.imgH = 0, 0
        self.createMode = createMode
        self.mode = self.CREATE
        self.scale = 1.0
        self.line = Shape()
        self.visible = {}
        self.shapes = []
        self.hVertex = None
        self.hEdge = None
        self.hShape = None
        self.prevhVertex = None
        self.prevhEdge = None
        self.prevhShape = None
        self.Vertex = None
        self.Edge = None
        self.Shape = None
        self.movingShape = False
        self.selectedShapes = []  # save the selected shapes here
        self.epsilon = 10.0

        self.setMouseTracking(True)

    def drawing(self):
        return self.mode == self.CREATE

    def editing(self):
        return self.mode == self.EDIT

    def setImage(self, imgFile):
        self.image = cv2.imread(imgFile)
        img = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
        img = QImage(img[:], img.shape[1], img.shape[0], img.shape[1] * 3, QImage.Format_RGB888)
        self.imgW, self.imgH = img.width(), img.height()
        img = img.scaled(QSize(self.width(), self.height()), Qt.KeepAspectRatio)
        self.setPixmap(QPixmap.fromImage(img))
        
    def offsetToCenter(self):
        s = self.scale
        area = super().size()
        w, h = self.imgW * s, self.imgH * s
        aw, ah = area.width(), area.height()
        x = (aw - w) / (2 * s) if aw > w else 0
        y = (ah - h) / (2 * s) if ah > h else 0
        return QPoint(x, y)

    def transformPos(self, point):
        return point / self.scale - self.offsetToCenter()

    def finalise(self):
        assert self.current
        self.current.close()
        self.shapes.append(self.current)
        #self.storeShapes()
        self.current = None
        #self.setHiding(False)
        #self.newShape.emit()
        self.update()

    def selectedVertex(self):
        return self.hVertex is not None

    def calculateOffsets(self, shape, point):
        rect = shape.boundingRect()
        x1 = rect.x() - point.x()
        y1 = rect.y() - point.y()
        x2 = (rect.x() + rect.width() - 1) - point.x()
        y2 = (rect.y() + rect.height() - 1) - point.y()
        self.offsets = QPoint(x1, y1), QPoint(x2, y2)

    def selectShapePoint(self, point):
        if self.selectedVertex():  # A vertex is marked for selection.
            index, shape = self.hVertex, self.hShape
            shape.highlightVertex(index, shape.MOVE_VERTEX)
        else:
            for shape in reversed(self.shapes):
                if shape.containsPoint(point):
                    self.calculateOffsets(shape, point)
                    #self.setHiding()
                    self.selectionChanged.emit([shape])
                    return
        #self.deSelectShape()

    def unHighlight(self):
        if self.hShape:
            self.hShape.highlightClear()
            self.update()
        self.prevhShape = self.hShape
        self.prevhVertex = self.hVertex
        self.prevhEdge = self.hEdge
        self.hShape = self.hVertex = self.hEdge = None

    def mousePressEvent(self, event):
        pos = self.transformPos(event.localPos())
        if event.button() == Qt.LeftButton:
            self.mode = self.CREATE
            if self.drawing():
                #if self.current:
                #    if self.createMode in ['rectangle', 'circle', 'line']:
                #        assert len(self.current.points) == 1
                #        self.current.points = self.line.points
                #        self.finalise()
                #    pass
                #else:
                self.current = Shape(shape_type=self.createMode)
                self.current.addPoint(pos)
                if self.createMode == 'circle':
                    self.current.shape_type = 'circle'
                self.line.points = [pos, pos]
                self.update()
            else:
                self.selectShapePoint(pos)
        elif event.button() == Qt.RightButton:
            #self.mode = self.EDIT
            print('MousePressEvent RightButton')

    def mouseMoveEvent(self, event):
        pos = self.transformPos(event.localPos())

        #print('{}, {}'.format(pos, self.mode))
        if self.drawing():
            if not self.current:
                return

            self.line.shape_type = self.createMode
            if self.createMode == 'rectangle':
                self.line.points = [self.current[0], pos]
                self.line.close()
            self.repaint()
            self.current.highlightClear()
            return
        else:
            for shape in self.shapes:
                # Look for a nearby vertex to highlight. If that fails,
                # check if we happen to be inside a shape.
                index = shape.nearestVertex(pos, self.epsilon / self.scale)
                index_edge = shape.nearestEdge(pos, self.epsilon / self.scale)
                print('{}, {}'.format(index, index_edge))
                if index is not None:
                    if self.selectedVertex():
                        self.hShape.highlightClear()
                    self.prevhVertex = self.hVertex = index
                    self.prevhShape = self.hShape = shape
                    self.prevhEdge = self.hEdge = index_edge
                    shape.highlightVertex(index, shape.MOVE_VERTEX)
                    self.update()
                    break
                elif shape.containsPoint(pos):
                    if self.selectedVertex():
                        self.hShape.highlightClear()
                    self.prevhVertex = self.hVertex
                    self.hVertex = None
                    self.prevhShape = self.hShape = shape
                    self.prevhEdge = self.hEdge = index_edge
                    self.update()
                    break
            else:  # Nothing found, clear highlights, reset state.
                self.unHighlight()
            return


        if Qt.LeftButton & event.buttons():
            if self.selectedVertex():
                self.boundedMoveVertex(pos)
                self.repaint()
                self.movingShape = True
            elif self.selectedShapes:# and self.prevPoint:
                #self.overrideCursor(CURSOR_MOVE)
                #self.boundedMoveShapes(self.selectedShapes, pos)
                #self.repaint()
                self.movingShape = True
            return

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.mode = self.EDIT
            print(len(self.shapes))
            self.finalise()
            pass

    def paintEvent(self, event):
        super().paintEvent(event)
        qp = QPainter()
        qp.begin(self)

        #for shape in self.shapes:
        #    if shape.selected and self.isVisible(shape):
        #        shape.fill = shape.selected or shape == self.hShape
        #        shape.paint(qp)
        if self.current:
            self.current.paint(qp)
            self.line.paint(qp)

        if self.shapes:
            for s in self.shapes:
                s.paint(qp)

        qp.end()
        
    def boundedMoveVertex(self, pos):
        index, shape = self.hVertex, self.hShape
        point = shape[index]
        #if self.outOfPixmap(pos):
        #    pos = self.intersectionPoint(point, pos)
        shape.moveVertexBy(index, pos - point)

    def boundedMoveShapes(self, shapes, pos):
        #if self.outOfPixmap(pos):
        #    return False  # No need to move
        o1 = pos + self.offsets[0]
        #if self.outOfPixmap(o1):
        #    pos -= QtCore.QPoint(min(0, o1.x()), min(0, o1.y()))
        o2 = pos + self.offsets[1]
        #if self.outOfPixmap(o2):
        #    pos += QtCore.QPoint(min(0, self.pixmap.width() - o2.x()),
        #                         min(0, self.pixmap.height() - o2.y()))
        # XXX: The next line tracks the new position of the cursor
        # relative to the shape, but also results in making it
        # a bit "shaky" when nearing the border and allows it to
        # go outside of the shape's area for some reason.
        # self.calculateOffsets(self.selectedShapes, pos)
        dp = pos - self.prevPoint
        if dp:
            for shape in shapes:
                shape.moveBy(dp)
            self.prevPoint = pos
            return True
        return False

    def resizeEvent(self, event):
        self.resized.emit()

class Widget(QWidget):
    def __init__(self, mode='rectangle', file=''):
        super().__init__()
        self.mode = mode
        self.file = file
        self.setupUi()
        self.setMouseTracking(True)

    def setupUi(self):
        self.setObjectName("Form")
        self.resize(954, 650)
        self.gridLayout = QGridLayout(self)
        self.gridLayout.setObjectName("gridLayout")

        self.lbImg = ImageLabel(self, self.mode)
        self.lbImg.setMinimumSize(QSize(622, 568))
        self.lbImg.setObjectName("lbImg")
        self.lbImg.setAlignment(Qt.AlignCenter)
        self.lbImg.setImage(self.file)
        self.gridLayout.addWidget(self.lbImg, 0, 0, 1, 1)

        self.retranslateUi()
        QMetaObject.connectSlotsByName(self)

    def retranslateUi(self):
        _translate = QCoreApplication.translate
        self.setWindowTitle(_translate(f'LabelImage {self.file}', f'标注图片 {self.file}'))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-m', '--mode', default='rectangle', help='Identify type: [rectangle, circle, line].')
    parser.add_argument('-f', '--file', default='', help='Default picture that will be show.')
    args = parser.parse_args()
    mode = args.mode
    file = args.file

    app = QApplication(sys.argv)
    w = Widget(mode=mode, file=file)
    w.show()
    sys.exit(app.exec())

